pkg/wshrpc/wshserver/wshserver.go (1)

702-734: Add logging and use constants for WSL prefix.

The implementation looks good but could benefit from the following improvements:

1. Add logging for better observability
2. Use a constant for the "wsl://" prefix

Apply this diff to improve the implementation:

+const wslPrefix = "wsl://"
+
 func (ws *WshServer) ConnUpdateWshCommand(ctx context.Context, updateInfo wshrpc.RemoteInfo) (bool, error) {
 	connName := updateInfo.ConnName
+	log.Printf("checking if wsh update is needed for %s", connName)
 	// check version number, return now if update not necessary
 	upToDate, _, err := conncontroller.IsWshVersionUpToDate(updateInfo.ClientVersion)
 	if err != nil {
 		return false, fmt.Errorf("unable to compare wsh version: %w", err)
 	}
 	if upToDate {
+		log.Printf("wsh is up to date for %s", connName)
 		// no need to update
 		return false, nil
 	}
+	log.Printf("wsh update needed for %s", connName)

 	// todo: need to add user input code here for validation

-	if strings.HasPrefix(connName, "wsl://") {
+	if strings.HasPrefix(connName, wslPrefix) {
 		return false, fmt.Errorf("cannot use this command for wsl connections")
 	}
 	connOpts, err := remote.ParseOpts(connName)
 	if err != nil {
 		return false, fmt.Errorf("error parsing connection name: %w", err)
 	}
 	conn := conncontroller.GetConn(ctx, connOpts, false, &wshrpc.ConnKeywords{})
 	if conn == nil {
 		return false, fmt.Errorf("connection not found: %s", connName)
 	}
 	err = conn.UpdateWsh(ctx, connName, &conncontroller.WshInstallOpts{Force: true, NoUserPrompt: true})
 	if err != nil {
 		return false, fmt.Errorf("update failed: %w", err)
 	}
+	log.Printf("successfully updated wsh for %s", connName)

 	// todo: need to add code for modifying configs?
 	return true, nil
 }

Track the TODO items.

There are two TODO items that need to be addressed:

1. Add user input code for validation
2. Add code for modifying configs

Would you like me to open GitHub issues to track these tasks?

pkg/wshutil/wshutil.go (1)

541-550: Consider adding error handling for environment variable.

The function assumes the SHELL environment variable exists. While this is typically true on Unix-like systems, it might not be set in all environments.

Consider this improvement:

func GetInfo(rpcContext wshrpc.RpcContext) wshrpc.RemoteInfo {
+    shell := os.Getenv("SHELL")
+    if shell == "" {
+        shell = "/bin/sh"  // provide a sensible default
+    }
    return wshrpc.RemoteInfo{
        ConnName:      rpcContext.Conn,
        ClientArch:    runtime.GOARCH,
        ClientOs:      runtime.GOOS,
        ClientVersion: wavebase.WaveVersion,
-        Shell:         os.Getenv("SHELL"),
+        Shell:         shell,
    }
}

Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro

pkg/remote/conncontroller/conncontroller.go (1)

380-393: Improve error handling and logging.

The function has the following suggestions:

1. Add more descriptive error messages
2. Use the connection's Infof method for consistent logging
3. Consider validating the RemoteInfo fields before use

Apply this diff to improve the implementation:

 func (conn *SSHConn) UpdateWsh(ctx context.Context, clientDisplayName string, remoteInfo *wshrpc.RemoteInfo) error {
-	log.Printf("attempting to update wsh for connection %v", remoteInfo)
+	conn.Infof(ctx, "attempting to update wsh for connection %s (os:%s arch:%s version:%s)\n",
+		remoteInfo.ConnName, remoteInfo.ClientOs, remoteInfo.ClientArch, remoteInfo.ClientVersion)
+
+	if remoteInfo.ClientOs == "" || remoteInfo.ClientArch == "" {
+		return fmt.Errorf("invalid remote info: missing os or arch information")
+	}
+
 	client := conn.GetClient()
 	if client == nil {
-		return fmt.Errorf("client is nil")
+		return fmt.Errorf("cannot update wsh: ssh client is not connected")
 	}
+
 	err := remote.CpWshToRemote(ctx, client, remoteInfo.ClientOs, remoteInfo.ClientArch)
 	if err != nil {
 		return fmt.Errorf("error installing wsh to remote: %w", err)
 	}
-	log.Printf("successfully installed wsh on %s\n", conn.GetName())
+	conn.Infof(ctx, "successfully updated wsh on %s\n", conn.GetName())
 	return nil
 }

pkg/wshrpc/wshserver/wshserver.go (2)

715-715: Track TODOs with GitHub issues.

There are two TODO comments that need to be tracked:

1. Add user input code for validation
2. Add code for modifying configs

Would you like me to create GitHub issues to track these TODOs?

Also applies to: 733-733

---

702-735: Improve error handling and add validation.

The implementation could be improved with:

1. More descriptive error messages
2. Validation of RemoteInfo fields
3. Consistent logging

Apply this diff to improve the implementation:

 func (ws *WshServer) ConnUpdateWshCommand(ctx context.Context, remoteInfo wshrpc.RemoteInfo) (bool, error) {
 	connName := remoteInfo.ConnName
+	if connName == "" {
+		return false, fmt.Errorf("invalid remote info: missing connection name")
+	}
+
+	log.Printf("checking wsh version for connection %s (current: %s)", connName, remoteInfo.ClientVersion)
+
 	// check version number, return now if update not necessary
 	upToDate, _, err := conncontroller.IsWshVersionUpToDate(remoteInfo.ClientVersion)
 	if err != nil {
 		return false, fmt.Errorf("unable to compare wsh version: %w", err)
 	}
 	if upToDate {
 		// no need to update
-		log.Printf("wsh update is not needed for connection %s", connName)
+		log.Printf("wsh is already up to date for connection %s", connName)
 		return false, nil
 	}

 	// todo: need to add user input code here for validation

 	if strings.HasPrefix(connName, "wsl://") {
-		return false, fmt.Errorf("cannot use this command for wsl connections")
+		return false, fmt.Errorf("wsh update is not supported for WSL connections")
 	}
 	connOpts, err := remote.ParseOpts(connName)
 	if err != nil {
 		return false, fmt.Errorf("error parsing connection name: %w", err)
 	}
 	conn := conncontroller.GetConn(ctx, connOpts, false, &wshrpc.ConnKeywords{})
 	if conn == nil {
 		return false, fmt.Errorf("connection not found: %s", connName)
 	}
+
+	log.Printf("attempting to update wsh for connection %s", connName)
 	err = conn.UpdateWsh(ctx, connName, &remoteInfo)
 	if err != nil {
-		return false, fmt.Errorf("update failed: %w", err)
+		return false, fmt.Errorf("wsh update failed for connection %s: %w", connName, err)
 	}

 	// todo: need to add code for modifying configs?
+	log.Printf("successfully updated wsh for connection %s", connName)
 	return true, nil
 }

Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro

pkg/remote/conncontroller/conncontroller.go (1)

380-394: ⚠️ Potential issue

Add validation for remoteInfo fields in UpdateWsh

Ensure that remoteInfo.ClientOs and remoteInfo.ClientArch are not empty before proceeding, to prevent potential errors in remote.CpWshToRemote.

Apply this diff to add validation:

     client := conn.GetClient()
     if client == nil {
         return fmt.Errorf("cannot update wsh: ssh client is not connected")
     }
+    if remoteInfo.ClientOs == "" || remoteInfo.ClientArch == "" {
+        return fmt.Errorf("invalid remote info: missing ClientOs or ClientArch")
+    }
     err := remote.CpWshToRemote(ctx, client, remoteInfo.ClientOs, remoteInfo.ClientArch)
     if err != nil {
         return fmt.Errorf("error installing wsh to remote: %w", err)
     }

pkg/wshrpc/wshserver/wshserver.go (2)

719-720: Implement user input validation

The TODO comment indicates that user input validation needs to be added. This is essential to ensure secure and reliable operation.

Would you like me to implement the user input validation code, or open a new GitHub issue to track this task?

---

737-738: Implement configuration modification code

The TODO comment suggests that code for modifying configurations is needed. Implementing this will ensure that configurations are properly updated after the WSH update.

Would you like me to help implement the configuration modification code, or open a new GitHub issue to track this task?

pkg/remote/conncontroller/conncontroller.go (3)

Line range hint 231-242: Refactor IsWshVersionUpToDate to accept version strings directly

The function IsWshVersionUpToDate currently expects a version line in the format "wsh v0.10.4", which may not be flexible. Consider refactoring the function to accept version strings directly, improving reusability.

Apply this diff to modify the function:

 func IsWshVersionUpToDate(clientVersion string) (bool, string, error) {
-    wshVersionLine = strings.TrimSpace(wshVersionLine)
-    if wshVersionLine == "not-installed" {
+    clientVersion = strings.TrimSpace(clientVersion)
+    if clientVersion == "not-installed" || clientVersion == "" {
         return false, "", nil
     }
-    parts := strings.Fields(wshVersionLine)
-    if len(parts) != 2 {
-        return false, "", fmt.Errorf("unexpected version format: %s", wshVersionLine)
-    }
-    clientVersion := parts[1]
     expectedVersion := fmt.Sprintf("v%s", wavebase.WaveVersion)
     if semver.Compare(clientVersion, expectedVersion) < 0 {
         return false, clientVersion, nil

Then, adjust the calls to this function accordingly.

---

Line range hint 293-299: Adjust call to IsWshVersionUpToDate

After refactoring IsWshVersionUpToDate, update the call in StartConnServer to pass the client version directly.

Apply this diff:

     conn.Infof(ctx, "got connserver version: %s\n", strings.TrimSpace(versionLine))
-    isUpToDate, clientVersion, err := IsWshVersionUpToDate(versionLine)
+    parts := strings.Fields(versionLine)
+    if len(parts) != 2 {
+        return false, "", fmt.Errorf("unexpected version format: %s", versionLine)
+    }
+    clientVersion := parts[1]
+    isUpToDate, clientVersion, err := IsWshVersionUpToDate(clientVersion)
     if err != nil {
         sshSession.Close()
         return false, "", fmt.Errorf("error checking wsh version: %w", err)

---

293-294: Handle potential nil clientVersion after parsing

After extracting clientVersion, ensure it is not empty before proceeding with the version comparison to prevent potential panics.

Apply this diff:

     clientVersion := parts[1]
+    if clientVersion == "" {
+        return false, "", fmt.Errorf("client version is empty")
+    }
     isUpToDate, clientVersion, err := IsWshVersionUpToDate(clientVersion)

Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro